{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6228ae7a",
   "metadata": {},
   "source": [
    "# Important Note:\n",
    "\n",
    "### I strongly suggest you follow my step-by-step video tutorial. The video will not only walk you through each step of the process but also includes a real-time execution using Diablo 2 as an example.\n",
    "\n",
    "### Here is the link to the video: \n",
    "https://www.youtube.com/watch?v=RSXgyDf2ALo\n",
    "[![Watch the video](https://img.youtube.com/vi/RSXgyDf2ALo/maxresdefault.jpg)](https://youtu.be/RSXgyDf2ALo)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c75577ac",
   "metadata": {},
   "source": [
    "# **Instructions**:\n",
    "\n",
    "### 1 - Open your game that you want to perform detections\n",
    "### 2 - In the game window, get the name of it's title bar \n",
    "### 3 - Update the variable \"window_name\" with the game title bar name\n",
    "### 4 - Run all cells to start generating your dataset in the folder images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91800ec7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import win32gui, win32ui, win32con\n",
    "from PIL import Image\n",
    "from time import sleep\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f9c9c67",
   "metadata": {},
   "outputs": [],
   "source": [
    "class WindowCapture:\n",
    "    w = 0\n",
    "    h = 0\n",
    "    hwnd = None\n",
    "\n",
    "    def __init__(self, window_name):\n",
    "        self.hwnd = win32gui.FindWindow(None, window_name)\n",
    "        if not self.hwnd:\n",
    "            raise Exception('Window not found: {}'.format(window_name))\n",
    "\n",
    "        window_rect = win32gui.GetWindowRect(self.hwnd)\n",
    "        self.w = window_rect[2] - window_rect[0]\n",
    "        self.h = window_rect[3] - window_rect[1]\n",
    "\n",
    "        border_pixels = 8\n",
    "        titlebar_pixels = 30\n",
    "        self.w = self.w - (border_pixels * 2)\n",
    "        self.h = self.h - titlebar_pixels - border_pixels\n",
    "        self.cropped_x = border_pixels\n",
    "        self.cropped_y = titlebar_pixels\n",
    "\n",
    "    def get_screenshot(self):\n",
    "        wDC = win32gui.GetWindowDC(self.hwnd)\n",
    "        dcObj = win32ui.CreateDCFromHandle(wDC)\n",
    "        cDC = dcObj.CreateCompatibleDC()\n",
    "        dataBitMap = win32ui.CreateBitmap()\n",
    "        dataBitMap.CreateCompatibleBitmap(dcObj, self.w, self.h)\n",
    "        cDC.SelectObject(dataBitMap)\n",
    "        cDC.BitBlt((0, 0), (self.w, self.h), dcObj, (self.cropped_x, self.cropped_y), win32con.SRCCOPY)\n",
    "\n",
    "        signedIntsArray = dataBitMap.GetBitmapBits(True)\n",
    "        img = np.fromstring(signedIntsArray, dtype='uint8')\n",
    "        img.shape = (self.h, self.w, 4)\n",
    "\n",
    "        dcObj.DeleteDC()\n",
    "        cDC.DeleteDC()\n",
    "        win32gui.ReleaseDC(self.hwnd, wDC)\n",
    "        win32gui.DeleteObject(dataBitMap.GetHandle())\n",
    "\n",
    "        img = img[...,:3]\n",
    "        img = np.ascontiguousarray(img) \n",
    "            \n",
    "        return img\n",
    "\n",
    "    def generate_image_dataset(self):\n",
    "        if not os.path.exists(\"images\"):\n",
    "            os.mkdir(\"images\")\n",
    "        while(True):\n",
    "            img = self.get_screenshot()\n",
    "            im = Image.fromarray(img[..., [2, 1, 0]])\n",
    "            im.save(f\"./images/img_{len(os.listdir('images'))}.jpg\")\n",
    "            sleep(0.3)\n",
    "    \n",
    "    def get_window_size(self):\n",
    "        return (self.w, self.h)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5507d47c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Execute this cell to generate a dataset of images for the specified window.\n",
    "\n",
    "window_name = \"YOUR_GAME_WINDOW_NAME_HERE\"\n",
    "\n",
    "wincap = WindowCapture(window_name)\n",
    "wincap.generate_image_dataset()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "140eee80",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
